package net.cyue.ort.llm.config;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;

import net.cyue.ort.DataType;
import net.cyue.util.StreamUtil;

import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Map;

public class ModelConfig {
    public final DataType dtype;
    public final int numHiddenLayers;
    public final int numAttentionHeads;
    public final int numKeyValueHeads;
    public final int headDim;
    public final Long bosTokenID;
    public final Long eosTokenID;
    public boolean useCache;

    public ModelConfig(
        DataType dtype,
        int numHiddenLayers,
        int numAttentionHeads,
        int numKeyValueHeads,
        int hiddenSize,
        Long bosTokenID,
        Long eosTokenID,
        boolean useCache
    ) {
        this.dtype = dtype;
        this.numHiddenLayers = numHiddenLayers;
        this.numAttentionHeads = numAttentionHeads;
        this.numKeyValueHeads = numKeyValueHeads;
        this.headDim = hiddenSize / numAttentionHeads;
        this.bosTokenID = bosTokenID;
        this.eosTokenID = eosTokenID;
        this.useCache = useCache;
    }

    public static ModelConfig fromFile(String configPath) throws IOException {
        return ModelConfig.fromInputStream(new FileInputStream(configPath));
    }
    public static ModelConfig fromInputStream(InputStream stream) throws IOException {
        return ModelConfig.fromContent(StreamUtil.readAsString(stream));
    }
    public static ModelConfig fromContent(String content) throws JsonProcessingException {
        ObjectMapper objectMapper = new ObjectMapper();
        Map<String, Object> configMap = null;
        configMap = objectMapper.readValue(
            content,
            new TypeReference<>() {}
        );

        // 读取模型配置
        String dtypeStr = (String) configMap.get("torch_dtype");
        if (dtypeStr == null) {
            dtypeStr = "float32";
        }
        DataType dtype = DataType.fromString(dtypeStr);
        int numHiddenLayers = (int) configMap.get("num_hidden_layers");
        int numAttentionHeads = (int) configMap.get("num_attention_heads");
        int numKeyValueHeads = (int) configMap.get("num_key_value_heads");
        int hiddenSize = (int) configMap.get("hidden_size");
        Integer bosTokenID = (Integer) configMap.get("bos_token_id");
        Integer eosTokenID = (Integer) configMap.get("eos_token_id");
        boolean useCache = (boolean) configMap.get("use_cache");
        return new ModelConfig(
            dtype,
            numHiddenLayers,
            numAttentionHeads,
            numKeyValueHeads,
            hiddenSize,
            bosTokenID == null ? null : (long) bosTokenID,
            eosTokenID == null ? null : (long) eosTokenID,
            useCache
        );
    }


    @Override
    public String toString() {
        return "ModelConfig {" +
            "\n\tdtype=" + dtype +
            ",\n\tnumHiddenLayers=" + numHiddenLayers +
            ",\n\tnumAttentionHeads=" + numAttentionHeads +
            ",\n\tnumKeyValueHeads=" + numKeyValueHeads +
            ",\n\theadDim=" + headDim +
            ",\n\tbosTokenID=" + bosTokenID +
            ",\n\teosTokenID=" + eosTokenID +
            ",\n\tuseCache="+ useCache +
            "\n}";
    }


    public ModelConfig setUseCache(boolean useCache) {
        this.useCache = useCache;
        return this;
    }


    public DataType getDtype() {
        return this.dtype;
    }

    public int getNumHiddenLayers() {
        return this.numHiddenLayers;
    }

    public int getNumAttentionHeads() {
        return this.numAttentionHeads;
    }

    public int getNumKeyValueHeads() {
        return this.numKeyValueHeads;
    }

    public int getHeadDim() {
        return this.headDim;
    }

    public Long getBosTokenID() {
        return this.bosTokenID;
    }

    public Long getEosTokenID() {
        return this.eosTokenID;
    }

    public boolean isUseCache() {
        return this.useCache;
    }


    // 常用模型配置
    public static ModelConfig SmolLM2_135M() {
        return new ModelConfig(
            DataType.FLOAT16,
            30,
            9,
            3,
            576,
            1L,
            2L,
            true
        );
    }
}
